/**
 * Copyright (c) 2012-2016, www.tinygroup.org (luo_guo@icloud.com).
 * <p>
 * Licensed under the GPL, Version 3.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 * <p>
 * http://www.gnu.org/licenses/gpl.html
 * <p>
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.tinygroup.alipayxmlsignature;

import org.tinygroup.xmlsignature.config.XmlSignatureConfig;
import org.tinygroup.xmlsignature.impl.DsigXmlSignatureProcessor;
import org.w3c.dom.Document;
import org.w3c.dom.Element;
import org.w3c.dom.Node;
import org.w3c.dom.NodeList;

import javax.xml.crypto.dom.DOMStructure;
import javax.xml.crypto.dsig.*;
import javax.xml.crypto.dsig.dom.DOMSignContext;
import javax.xml.crypto.dsig.dom.DOMValidateContext;
import javax.xml.crypto.dsig.keyinfo.KeyInfo;
import javax.xml.crypto.dsig.spec.C14NMethodParameterSpec;
import javax.xml.crypto.dsig.spec.TransformParameterSpec;
import java.io.OutputStream;
import java.security.KeyPair;
import java.util.Collections;
import java.util.List;

/**
 * 支付宝XML数字签名(Finance报文)
 * @author yancheng11334
 *
 */
public class FinanceXmlSignatureProcessor extends DsigXmlSignatureProcessor {

    protected List<Reference> createReference(Document doc, XmlSignatureConfig config) throws Exception {
        Transform envelopedTransform = xmlSignatureFactory.newTransform(Transform.ENVELOPED,
                (TransformParameterSpec) null);
        DigestMethod sha1DigMethod = xmlSignatureFactory.newDigestMethod(DigestMethod.SHA1, null);

        String referName = "#" + getReferName(doc);
        Reference ref = xmlSignatureFactory.newReference(referName, sha1DigMethod, Collections.singletonList(envelopedTransform), null, null);
        return Collections.singletonList(ref);
    }

    private String getReferName(Document doc) throws Exception {
        Element messageElement = (Element) getSignatureXmlNode(doc);
        NodeList nodelist = messageElement.getChildNodes();
        for (int i = 0; i < nodelist.getLength(); i++) {
            Node node = nodelist.item(i);
            if (node instanceof Element) {
                Element ele = (Element) node;
                if (ele.getAttribute("id") != null) {
                    return ele.getAttribute("id");
                }
            }
        }
        return null;
    }

    protected Node getSignatureXmlNode(Document doc)
            throws Exception {
        return doc.getElementsByTagName("Message").item(0);
    }

    protected CanonicalizationMethod createCanonicalizationMethod(
            XmlSignatureConfig config) throws Exception {
        CanonicalizationMethod cmethod = xmlSignatureFactory.newCanonicalizationMethod(CanonicalizationMethod.INCLUSIVE, (C14NMethodParameterSpec) null);
        return cmethod;
    }

    protected SignatureMethod createSignatureMethod(XmlSignatureConfig config)
            throws Exception {
        SignatureMethod smethod = xmlSignatureFactory.newSignatureMethod(SignatureMethod.RSA_SHA1, null);
        return smethod;
    }

    //支付宝XML数字签名规范：签名元素不包含keyInfo节点
    protected KeyInfo createKeyInfo(XmlSignatureConfig config) throws Exception {
        return null;
    }

    protected void createXmlSignature(XmlSignatureConfig config, Document doc,
                                      OutputStream output) throws XMLSignatureException {
        try {
            CanonicalizationMethod cm = createCanonicalizationMethod(config);
            SignatureMethod sm = createSignatureMethod(config);
            List<Reference> references = createReference(doc, config);
            SignedInfo si = createSignedInfo(cm, sm, references);
            KeyInfo ki = createKeyInfo(config);
            XMLSignature signature = createXMLSignature(si, ki);

            Node messageNode = getSignatureXmlNode(doc);
            KeyPair keyPair = getXmlSignatureManager().getKeyPair(config.getUserId());
            DOMSignContext signContext = new DOMSignContext(keyPair.getPrivate(), messageNode);
            //String referName = getReferName(doc);
            //Element element = (Element)doc.getElementsByTagName(referName).item(0);
            //signContext.setIdAttributeNS(element, null, "id");
            signature.sign(signContext);

            transform(doc, output);
        } catch (Exception e) {
            throw new XMLSignatureException("生成支付宝XML数字签名失败", e);
        }
    }

    protected boolean validateXmlSignature(XmlSignatureConfig config,
                                           Document doc) throws XMLSignatureException {
        try {
            // 查找签名元素
            NodeList nl = doc.getElementsByTagNameNS(XMLSignature.XMLNS,
                    "Signature");
            if (nl.getLength() == 0) {
                throw new Exception("没有找到<Signature>元素");
            }
            Node signatureNode = nl.item(0);
            XMLSignature signature = xmlSignatureFactory.unmarshalXMLSignature(new DOMStructure(signatureNode));

            KeyPair keyPair = getXmlSignatureManager().getKeyPair(config.getUserId());
            DOMValidateContext valCtx = new DOMValidateContext(keyPair.getPublic(),
                    signatureNode);
            return signature.validate(valCtx);
        } catch (Exception e) {
            throw new XMLSignatureException("验证支付宝XML数字签名失败", e);
        }
    }


}
